{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "302c696f",
   "metadata": {},
   "outputs": [],
   "source": [
    "from numpy import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "86c1583d",
   "metadata": {},
   "outputs": [],
   "source": [
    "class HMM:\n",
    "    def __init__(self,A,B,pi):\n",
    "        self.A =A\n",
    "        self.B =B\n",
    "        self.pi =pi\n",
    "    def _forward(self,obs_seq):\n",
    "        # 取A = N x N\n",
    "        N = self.A.shape[0]\n",
    "        T = len(obs_seq)\n",
    "\n",
    "        F = zeros((N,T))\n",
    "\n",
    "        # alpha = pi*b\n",
    "        F[:,0] = self.pi *self.B[:,obs_seq[0]]\n",
    "\n",
    "        for t in range(1,T):\n",
    "            for n in range(N):\n",
    "                # 计算第t时，第n个状态的前向概率\n",
    "                F[n,t] = dot(F[:,t-1], (self.A[:,n])) * self.B[n, obs_seq[t]]\n",
    "        return F\n",
    "    def _backward(self,obs_seq):\n",
    "        N = self.A.shape[0]\n",
    "        T = len(obs_seq)\n",
    "\n",
    "        X = zeros((N,T))\n",
    "        # 表示X矩阵的最后一列\n",
    "        X[:,-1:] =1\n",
    "\n",
    "        for t in reversed(range(T-1)):\n",
    "            for n in range(N):\n",
    "                # 边权值为a_ji\n",
    "                X[n,t] = sum(X[:,t+1] * self.A[n,:] * self.B[:,obs_seq[t+1]])\n",
    "\n",
    "        return X\n",
    "    def viterbi(self,obs_seq):\n",
    "\n",
    "        N = self.A.shape[0]\n",
    "        T = len(obs_seq)\n",
    "\n",
    "        prev = zeros((T-1,N),dtype = int)\n",
    "\n",
    "        V = zeros((N,T))\n",
    "        V[:,0] = self.pi * self.B[:,obs_seq[0]]\n",
    "\n",
    "        for  t in range(1,T):\n",
    "            for n in range(N):\n",
    "                # 计算delta(j)*a_ji\n",
    "                seq_probs = V[:,t-1] * self.A[:,n] * self.B[n,obs_seq[t]]\n",
    "                # 记录最大状态转移过程\n",
    "                prev[t-1,n] = argmax(seq_probs)\n",
    "                V[n,t] = max(seq_probs)\n",
    "        return V,prev\n",
    "\n",
    "    def build_viterbi_path(self,prev,last_state):\n",
    "        \"\"\"\n",
    "        returns a state path ending in last_state in reverse order.\n",
    "        \"\"\"\n",
    "        T = len(prev)\n",
    "        yield(last_state)\n",
    "        # 从T-1开始，每次下降1\n",
    "        for i in range(T-1,-1,-1):\n",
    "            yield(prev[i,last_state])\n",
    "            last_state = prev[i,last_state]\n",
    "\n",
    "    def state_path(self,obs_seq):\n",
    "        V,prev = self.viterbi(obs_seq)\n",
    "\n",
    "        # build state path with greatest probability\n",
    "        last_state = argmax(V[:,-1])\n",
    "        path = list(self.build_viterbi_path(prev,last_state))\n",
    "\n",
    "        return V[last_state,-1],reversed(path)\n",
    "    \n",
    "    def baum_welch_train(self, observations, criterion=0.05):\n",
    "        n_states = self.A.shape[0]\n",
    "        # 观察序列的长度T\n",
    "        n_samples = len(observations)\n",
    "\n",
    "        done = False\n",
    "        while not done:\n",
    "            # alpha_t(i) = P(o_1,o_2,...,o_t,q_t = s_i | hmm)\n",
    "            # Initialize alpha\n",
    "            # 获得所有前向传播节点值 alpha_t(i)\n",
    "            alpha = self._forward(observations)\n",
    "\n",
    "            # beta_t(i) = P(o_t+1,o_t+2,...,o_T | q_t = s_i , hmm)\n",
    "            # Initialize beta\n",
    "            # 获得所有后向传播节点值 beta_t(i)\n",
    "            beta = self._backward(observations)\n",
    "\n",
    "            # 计算 xi_t(i,j) -> xi(i,j,t)\n",
    "            xi = np.zeros((n_states, n_states, n_samples - 1))\n",
    "            # 在每个时刻\n",
    "            for t in range(n_samples - 1):\n",
    "                # 计算P(O | hmm)\n",
    "                denom = sum(alpha[:, -1])\n",
    "                for i in range(n_states):\n",
    "                    # numer[1,:] = 行向量，alpha[i,t]=实数，slef.A[i,:] = 行向量\n",
    "                    # self.B[:,observations[t+1]].T = 行向量,beta[:,t+1].T = 行向量\n",
    "                    numer = alpha[i, t] * self.A[i, :] * self.B[:, observations[t + 1]].T * beta[:, t + 1].T\n",
    "                    xi[i, :, t] = numer / denom\n",
    "\n",
    "                # 计算gamma_t(i) 就是对j进行求和\n",
    "                gamma = np.sum(xi, axis=1)\n",
    "                # need final gamma elements for new B\n",
    "                prod = (alpha[:, n_samples - 1] * beta[:, n_samples - 1]).reshape((-1, 1))\n",
    "                # 合并T时刻的节点\n",
    "                gamma = np.hstack((gamma, prod / np.sum(prod)))\n",
    "                # 列向量\n",
    "                newpi = gamma[:, 0]\n",
    "                newA = np.sum(xi, 2) / np.sum(gamma[:, :-1], axis=1).reshape((-1, 1))\n",
    "                newB = np.copy(self.B)\n",
    "\n",
    "                # 观测状态数\n",
    "                num_levels = self.B.shape[1]\n",
    "                sumgamma = np.sum(gamma, axis=1)\n",
    "                for lev in range(num_levels):\n",
    "                    mask = observations == lev\n",
    "                    newB[:, lev] = np.sum(gamma[:, mask], axis=1) / sumgamma\n",
    "\n",
    "                if np.max(abs(self.pi - newpi)) < criterion and \\\n",
    "                                np.max(abs(self.A - newA)) < criterion and \\\n",
    "                                np.max(abs(self.B - newB)) < criterion:\n",
    "                    done = 1\n",
    "                self.A[:], self.B[:], self.pi[:] = newA, newB, newpi\n",
    "    def simulate(self, T):\n",
    "        def draw_from(probs):\n",
    "            # np.random.multinomial 为多项式分布，1为实验次数，类似于投掷一枚骰子，丢出去是几，probs每个点数的概率，均为1/6\n",
    "            # 给定行向量的概率，投掷次数为1次，寻找投掷的点数\n",
    "            return np.where(np.random.multinomial(1, probs) == 1)[0][0]\n",
    "\n",
    "        observations = np.zeros(T, dtype=int)\n",
    "        states = np.zeros(T, dtype=int)\n",
    "        states[0] = draw_from(self.pi)\n",
    "        observations[0] = draw_from(self.B[states[0], :])\n",
    "        for t in range(1, T):\n",
    "            states[t] = draw_from(self.A[states[t - 1], :])\n",
    "            observations[t] = draw_from(self.B[states[t], :])\n",
    "        return observations, states"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a6d635fd",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
